from toolz import curry

import tensorflow as tf
from hanser.transform import _image_dimensions, _get_lam, rand_bbox


def rand_mask(image, lam):
    n, h, w = _image_dimensions(image, 4)[:3]
    l, t, r, b = rand_bbox(h, w, lam)
    hi = tf.range(h)[None, :, None, None]
    mh = (hi >= t[:, None, None, None]) & (hi < b[:, None, None, None])
    wi = tf.range(w)[None, None, :, None]
    mw = (wi >= l[:, None, None, None]) & (wi < r[:, None, None, None])
    masks = tf.cast(tf.logical_not(mh & mw), image.dtype)
    return masks


@curry
def cutmix_batch(image, alpha, hard=True, **gen_lam_kwargs):
    n = _image_dimensions(image, 4)[0]
    lam_shape = (n,) if hard else (1,)
    lam = _get_lam(lam_shape, alpha, **gen_lam_kwargs)

    mask_fn = rand_mask

    masks = mask_fn(image, lam)

    indices = tf.random.shuffle(tf.range(n))
    image2 = tf.gather(image, indices)

    image = image * masks + image2 * (1. - masks)
    return image, image
